﻿#################################################
#												#
# Newton's Method								#
# A fairly incomplete/experimental script		#
# Either reads a function from user input		#
# and finds a zero, or (statically) tries to	#
# find lambda-star for a variety of constant	#
# vorticities									#
#												#
# Matt Ball (ball@brown.edu)					#
# 2008											#
#												#
#################################################

# TODO:
# - plotting
# - f in terms of an integral (eval numerically)
# - look into using a pre-made newton's method in scipy
# - \Gamma will be entered symbolically
# - the overall eqn will be in terms of an integral and \Gamma

from sympy import *
from math import *
#from pylab import *
#from pyx import *
import Gnuplot

def getP0():
	"""
	Read in p_0 from user input
	"""
	p0 = input("Enter a NEGATIVE value for p0: ")
	while (p0 >= 0):
		print "\tp0 must be strictly negative."
		p0 = input("\tEnter a NEGATIVE value for p0: ")
	return p0
	
def getVorticity():
	"""
	Read in the vorticity strength (a constant) from user input
	"""
	v = input("Enter the value of the (constant) vorticity: ")
	return v
	
def getX(f):
	"""
	Lets the user specify the dependent variable (so it doesn't *have* to be x)
	"""
	x = Symbol(raw_input("What is the dependent variable? "))
	symbols = f.atoms(Symbol)
	
	# TODO - update this
	# Make sure the equation contains only x, or is constant
	while x not in symbols:
		print "\tThat is not a variable in the equation you entered."
		x = Symbol(raw_input("\tWhat is the dependent variable? "))
	return x
	
def simpleNewtonsMethod(initial,stoppingDelta,relation):
	"""
	Runs the standard 1D Newton's Method algorithm
	"""
	
	# initial values
	current = initial
	next = relation(initial)
	
	# Newton's Method:
	while (abs(next-current) >= stoppingDelta):		
		if (abs(next) == oo) or (abs(current) == oo):
			break
		
		current = next
		next = relation(current)
	
	return next

def getBasics():
	"""
	Ask the user for some input and runs Newton's Method accordingly.
	"""
	
	g = Real(9.80665) # m/s/s, or 32.174 ft/s/s
	
	print "Assumptions: g=9.8 m/s/s, p0 is relative mass flux, v is vorticity."
	f = sympify(raw_input("Enter an equation to solve: 0 = ")) # Get f
	x = getX(f)
	
	fprime = f.diff(x)	# Compute f'(x)
	
	symbols = f.atoms(Symbol)
	
	# Get p0, if it was used
	p0 = Symbol('p0')
	if p0 in symbols:
		p0 = Real(getP0())
		
	# Get vorticity, if it was used
	v = Symbol('v')
	if v in symbols:
		v = Real(getVorticity())
	
	# Get x0, the initial guess
	x0 = float(raw_input("Enter the initial guess for f(x)=0: "))
	
	print "f(" + `x` + ") =",f
	print "f'(" + `x` + ") =",fprime
	
	# sympify the expression for (hopefully) faster evaluation
	newt = x - f/fprime
	newt = sympify(newt)
	# Convert to python lambda	
	newt = lambdify(x,newt.subs({'g':g,'p0':p0,'v':v}),"math")
	
	delta = float(raw_input("Stop when |x_(n+1) - x_n| < what value? "))
	
	result = simpleNewtonsMethod(x0,delta,newt)
	print result
	return result
	
def lambdaStar():
	"""
	Statically tries to approximate lambda star - not so good.
	"""

	print len(sys.argv)
	
	p0 = -2
	delta = 1e-12 # python seems to do weird stuff with exponents smaller than -10
	g = 9.816 # correct units?
	
	x = Symbol('x')
	f = sympify("tanh(((x**.5) - (x + 2*p0*gam)**.5)/gam) - (x/(g + gam*(x**.5)))")
	fprime = f.diff(x)
	newt = sympify(x - f/fprime)
	
	#z_min = 0.1
	#z_max = 10.
	
	print " ɣ\t|\t λ*"
	
	# non-functioning graphing code below
	#for gam in range(-10,-8):
		#def f(z):
		#	return lambdify(z,newt.subs({'g':g,'p0':p0,'gam':gam}))
		
		#graf = graph.graphxy(width=10, x=graph.axis.linear(min=z_min, max=z_max), y=graph.axis.linear(min=-5, max=5))
		#graf.plot(graph.data.paramfunction("x",z_min,z_max,"x, y = x, f(z)",context=locals()))
		#graf.writePDFfile("lambdaStar/gamma=" + str(gam) + ".pdf")
	
	# a quick guess from excel, reasonable on [-10,10]
	def guess(gam):
		return (4E-06)*(gam**6) - (2E-05)*(gam**5) - 0.0012*(gam**4) + 0.0037*(gam**3) + 0.2167*(gam**2) + 1.7957*gam + 6.4773
		
	print guess(10)
	
	for gam in range(-20,21): # since range won't take non-int step size
		gam = gam/2.
		temp = lambdify(x,newt.subs({'g':g,'p0':p0,'gam':gam}))
		
		try:
			result = simpleNewtonsMethod(guess(gam),delta,temp)
		except ValueError:
			result = nan
		print gam, "\t|\t", result

def main():
	
	funcmap = dict(simple=getBasics, lambdaStar=lambdaStar)
	
	def usage():
		print 'Usage: python %s [%s]' % (sys.argv[0], ' | '.join(funcmap.keys()))
	
	funckey = len(sys.argv) > 1 and sys.argv[1] or None
	func = funckey in funcmap and funcmap[funckey] or usage
	func()

if __name__ == '__main__':
	main()
